import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
from tqdm import tqdm
import os
from CNN import Net1, Net2, Net3, Net4

# 定义数据转换，将输入通道从三通道改为单通道
transform = transforms.Compose([
    transforms.ToTensor(),
])

# 加载FashionMNIST数据集
trainset = torchvision.datasets.FashionMNIST(root='./data', train=True,
                                              download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=64, shuffle=True, drop_last=True)

testset = torchvision.datasets.FashionMNIST(root='./data', train=False,
                                             download=True, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=2, shuffle=False)

# 创建保存模型权重的文件夹
if not os.path.exists('ckpt'):
    os.makedirs('ckpt')

# 定义网络
networks = {
    "Net1": Net1(),
    "Net2": Net2(),
    "Net3": Net3(),
    "Net4": Net4(),
}

# 加载已保存的权重（如果存在）
for name, net in networks.items():
    weight_path = os.path.join('ckpt', name + '.pth')
    if os.path.exists(weight_path):
        net.load_state_dict(torch.load(weight_path))


# 定义损失函数
criterion = nn.CrossEntropyLoss()

# 定义优化器
optimizers = {
    "Net1": optim.Adam(networks["Net1"].parameters(), lr=0.001),
    "Net2": optim.Adam(networks["Net2"].parameters(), lr=0.001),
    "Net3": optim.Adam(networks["Net3"].parameters(), lr=0.001),
    "Net4": optim.Adam(networks["Net4"].parameters(), lr=0.001)
}

epochs = 20
# 训练和测试每个网络
for name, net in networks.items():
    # print("Training", name)
    # optimizer = optimizers[name]
    # for epoch in range(epochs):  # 多次循环遍历数据集
    #     running_loss = 0.0
    #     # 使用tqdm创建进度条
    #     with tqdm(total=len(trainloader), desc=f'Epoch {epoch + 1}/{epochs}', unit='batch') as pbar:
    #         for i, data in enumerate(trainloader, 0):
    #             inputs, labels = data
    #             optimizer.zero_grad()
    #
    #             outputs = net(inputs)
    #             loss = criterion(outputs, labels)
    #             loss.backward()
    #             optimizer.step()
    #
    #             running_loss += loss.item()
    #
    #             # 更新进度条
    #             pbar.set_postfix({'loss': running_loss / (i + 1)})
    #             pbar.update(1)

    # print('Finished Training', name)

    # 保存模型权重
    # save_path = os.path.join('ckpt', name + '.pth')
    # torch.save(net.state_dict(), save_path)

    # 测试网络
    correct = 0
    total = 0
    with torch.no_grad():
        for data in testloader:
            images, labels = data
            outputs = net(images)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

    print(f'Accuracy of the network {name} on the 10000 test images: {100 * correct / total} %')
